# 我们的图形之所以与书上的不同，我发现主要和valfun_varlab.R函数中
# ct小于0时，如何设置值函数有关。改变该值图形波动非常大。但作者并未
# 给出相应处理。这里是比较疑惑的。

rm(list = ls())
library(rootSolve)
library(parallel)
library(reshape2)
library(ggplot2)
library(magrittr)
library(stringr)
library(foreach)
library(signal)
library(patchwork)
devtools::load_all()

# 参数设置
k0 <- seq(0.06,10,0.06)
vlast <- rep(0,length(k0))
beta <- .98
delta <-  .1
theta <- .36
numits <- 240
A <- 0.5

# 图形初始化
picdata <- data.frame(k0 = k0, v = vlast)

# 值函数迭代numits次。在每个k=(0.01，6.2)上以0作为初值开始迭代
v <- k_opt <- h_opt<- numeric(length(k0))
cl <- makeCluster(10)
doParallel::registerDoParallel(cl)
for (i in 1:numits){
  # 寻找每个k点的最优值函数, 非平行运算版本，注意注释掉cl等平行运算的行代码。
  # 本处代码主要用于理解
  # for (j in 1:length(k0)){
    # 多变量带约束优化
  #   ans <- constrOptim(f = valfun_varlab,theta = c(1,0.5), grad = NULL,ui = matrix(c(0,0,1,-1),ncol = 2),
  #       ci = matrix(c(0,-1),ncol = 1), kt = k0[j], beta = beta, theta_val = theta,
  #       delta = delta, k0 = k0, vlast = vlast, A = A, control = list(fnscale = -1))
  #   # ans <- optim(par = c(0.6,0.5),fn = valfun_varlab, lower = c(0,0),
  #   #                    upper = c(10, 1), kt = k0[j], beta = beta, theta_val = theta,
  #   #                    delta = delta, k0 = k0, vlast = vlast, A = A, method = "L-BFGS-B",control = list(fnscale = -1))
  #
  #   v[j] <- ans$value
  #   k_opt[j] <- ans$par[1]
  #   h_opt[j] <- ans$par[2]
  # }

  # 寻找每个k点的最优值函数, 平行运算版本
  ans <- foreach(i = k0,vlast = rep(list(vlast),length(k0)),
                 .export = c('valfun_varlab','k0','beta','theta','delta','A')) %dopar% {
    constrOptim(f = valfun_varlab,theta = c(1,0.5), grad = NULL,ui = matrix(c(0,0,1,-1),ncol = 2),
                ci = matrix(c(0,-1),ncol = 1), kt = i, beta = beta, theta_val = theta,
                delta = delta, k0 = k0, vlast = vlast, A = A, control = list(fnscale = -1))
  }
  v <- sapply(1:length(k0),function(i) ans[[i]]$value)
  k_opt <- sapply(1:length(k0),function(i) ans[[i]]$par[1])
  h_opt <- sapply(1:length(k0),function(i) ans[[i]]$par[2])

  # 每48次迭代存储一下
  if (i %% 48 == 0){
    print(i)
    picdata[,paste('v',as.character(i), sep = '')] <- v
    picdata[,paste('k_opt',as.character(i), sep = '')] <- k_opt
    picdata[,paste('h_opt',as.character(i), sep = '')] <- h_opt
  }

  # 替换上一次的值函数
  vlast <- v
}
stopCluster(cl)

# value function
ans <- melt(picdata, id.vars = 'k0')
ggplot(ans[str_detect(ans$variable,'^v[1-9]'),], aes(x = k0, y = value, color = variable)) +
  geom_line() + theme_bw()

# policy function
ggplot(ans[str_detect(ans$variable,'k_opt240'),], aes(x = k0, y = value, color = variable)) +
  geom_line() + labs(y = 'knext') + theme_bw()

ggplot(ans[str_detect(ans$variable,'h_opt240'),], aes(x = k0, y = value, color = variable)) +
  geom_line() + labs(y = 'h') + theme_bw()

#--------内生格点法--------
# 格点化状态变量k, 同时猜测控制变量c, h
h <- k <- seq(0.06,10,length.out = 80) %>% matrix(ncol = 1)
cc <- seq(0.01,0.9,length.out = 80) %>% matrix(ncol = 1)

findh <- function(h,parms = list(cc,A,theta,k)){
  parms$A*parms$cc/(1-h)-(1-parms$theta)*parms$k^parms$theta*h^(-parms$theta)
}

for (i in 1:100) {
  knext <- (1-delta) * k + k^theta * h[,ncol(h)]^(1-theta)-cc[,ncol(cc)]
  cnext <- interp1(k, cc[,ncol(cc)],knext, extrap = T)

  # updata c and h
  cc <- (cnext/(beta*(theta*k^(theta-1)*h[,ncol(h)]^(1-theta)+1-delta))) %>% cbind(cc,.)
  htemp <- numeric(nrow(h))
  for (j in 1:nrow(cc)) {
    htemp[j] <- uniroot.all(findh,c(0.1,10), parms = list(A = A, theta = theta, cc = cc[j,ncol(cc)],k = k[j]))[1]
  }
  h <- cbind(h,htemp)
}

picdata <- data.frame(cc = cc[,ncol(cc)], k = k, h = h[,ncol(h)])
picdata <- within(picdata, {knext <- (1-delta) * k + k^theta * h^(1-theta)-cc})

# policy function
p1 <- ggplot(picdata, aes(x = k, y = knext)) +
  geom_line() + theme_bw()
p2 <- ggplot(picdata, aes(x = k, y = cc)) +
  geom_line(color = I('red')) + theme_bw()
p3 <- ggplot(picdata, aes(x = k, y = h)) +
  geom_line(color = I('steelblue1')) + theme_bw()
p1 + p2 + p3
